import numpy as np
from skimage.measure import block_reduce
import matplotlib.pyplot as plt
import cv2
from skimage.measure import label, regionprops, perimeter
from skimage.morphology import binary_erosion, square
from skimage.filters import threshold_otsu
from skimage.color import label2rgb
from skimage.transform import resize
from skimage.filters import gaussian
import numpy as np
from scipy import ndimage

def mark_high_density_areas(image, block_size, density_threshold,color):
    # Reduce the image to the specified block size
    reduced_image = block_reduce(image, block_size, np.sum)
    # Create a mask for the high density areas
    density_mask = reduced_image > density_threshold
    # Copy the original image
    marked_image = image.copy()
    # Mark the high density areas

    density_mask = np.repeat(density_mask, block_size[0], axis=0)
    density_mask = np.repeat(density_mask, block_size[1], axis=1)
    marked_image[density_mask] = color
    return marked_image


def connecte_image(image, pixel_value, distance, sigma_factor = 3, minimum_area_factor = 100):
    height, width = image.shape
    # 首先将图像中的像素值大于x的区域设置为1，其余区域设置为0
    binary_image = np.where(image >= pixel_value, 1, 0)
    # 使用binary_closing()函数连通距离小于y的区域
    closed_image = ndimage.binary_closing(binary_image, iterations=distance)
    # perform gaussian smoothing on the edges
    # 对连通区域进行标记
    labeled_image, num_features = ndimage.label(closed_image)

    # 找到每个连通区域的像素和，并将像素和小于min_sum的连通区域标记为0
    for i in range(1, num_features + 1):
        feature_pixels = np.where(labeled_image == i)
        feature_sum = np.sum(image[feature_pixels])
        if feature_sum < (height * width) / minimum_area_factor:
            closed_image[feature_pixels] = 0

    smooth_image = gaussian(closed_image, sigma=distance / sigma_factor)
    return smooth_image

def overlay_mask(rgb_img: np.ndarray, gray_img: np.ndarray, alpha: 0.5) -> np.ndarray:
    # 将灰度图中像素值大于零的区域转化为红色区域
    mask = gray_img > 0.008
    color_img = np.zeros_like(gray_img, dtype=np.uint8)
    color_img[mask] = 255

    # 重复3遍
    color_img = np.dstack((color_img, color_img, color_img))

    # 将颜色改为红色
    color_img[:, :, 0] = 0
    color_img[:, :, 1] = 0
    # 将红色标记区域叠加到RGB图片上
    result = cv2.addWeighted(rgb_img, 0.5, color_img, alpha, 0)
    return result


def get_warning_map(img, density_map, h=32, density_threshold=0.01, color=0.5, distance=18, sigma_factor = 3, minimum_area_factor = 100):
    block_size = (int(img.shape[0] / h), int(img.shape[1] / h))
    mark_map = mark_high_density_areas(density_map, block_size, density_threshold, color)
    connected_map = connecte_image(mark_map, color, distance, sigma_factor, minimum_area_factor)
    overlay_mask_map = overlay_mask(img, connected_map, alpha=0.4)
    return overlay_mask_map


# test code
if __name__=="__main__":
    # Example usage
    binary_value_map = np.random.randint(0, 3, (10, 10))
    marked_image = mark_high_density_areas(binary_value_map, (2, 2), 5,255)

    fig = plt.figure()

    plt.subplot(1, 2, 1)
    plt.imshow(binary_value_map)
    plt.title('binary_value_map')

    plt.subplot(1, 2, 2)
    plt.imshow(marked_image, cmap='gray')
    plt.title('marked_image')
